import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.neural_network import MLPClassifier
from imblearn.over_sampling import SMOTE
import pickle


def read_total_data(DataFrame):
    feature = []
    label = []
    for row in range(len(DataFrame)):
        temp_feature = []
        for col in range(len(DataFrame.columns) - 1):
            temp_feature.append(DataFrame[DataFrame.columns[col]][row])
        feature.append(temp_feature)
        label.append(DataFrame[DataFrame.columns[len(DataFrame.columns) - 1]][row])
    return feature, label


data = pd.read_csv("last_encode.csv")

total_feature, total_label = read_total_data(data)

total_feature_smote, total_label_smote = SMOTE(random_state=0).fit_resample(total_feature, total_label)

X_train, X_test, Y_train, Y_test = train_test_split(total_feature_smote,
                                                    total_label_smote,
                                                    random_state=0,
                                                    train_size=0.7,
                                                    test_size=0.3,
                                                    )

ANN = MLPClassifier(random_state=0).fit(X_train, Y_train)
print(ANN.score(X_test, Y_test))

# 写出去
with open("ANN.pickle", "wb") as fw:
    pickle.dump(ANN, fw)

# 读进来
'''with open('LR.pickle', 'rb') as fr:
    new_LR = pickle.load(fr)'''
